package midend.IR_Optimizer;

import midend.IR_Instruction.*;
import midend.IR_Module;
import midend.IR_Type.*;
import midend.IR_Value.*;

import java.lang.reflect.Array;
import java.util.*;

import static midend.IR_Type.IR_Const.*;

public class Branch_Eli implements IR_Opt {
    private IR_Module module;
    //private Map<IR_BaseBlock,IR_BaseBlock> bb2prebb;

    public Branch_Eli(IR_Module module) {
        this.module = module;
        module.buildCFG();
    }

    public void run() {
        for(IR_FuncBlock funcBlock : module.getFuncBlocks()) {
            Optimize(funcBlock);
        }
        module.buildCFG();
    }
    public void Optimize(IR_FuncBlock ir_funcBlock) {
        //todo 把 constmap转化为value to value 的删除 move 或者 zext 以及 cmp 常量的指令
        Map<IR_Value, IR_Value> constmap = new HashMap<>();
        boolean changed = true;
        while (changed) {
            changed = false;
            //a ,b  b前面插入a
            Map<Instruction, Instruction> insertmap = new HashMap<>();
            Set<Instruction> deletelist = new HashSet<>();
            Map<IR_BaseBlock,Integer> deleteblocklist = new HashMap<>();
            //每次遍历完所有已知常量 完成替换
            for (IR_BaseBlock bb : ir_funcBlock.get_bb_List()) {
                for (Instruction instruction : bb.get_Instruction_List()) {
                    if (instruction instanceof Cond_Instruction condInstruction) {
                        //System.out.println(condInstruction);
                        if(constmap.containsKey(condInstruction.get_op_index(1))) {
                            IR_Value val = constmap.get(condInstruction.get_op_index(1));
                            //System.out.println(val);
                            if (val instanceof IR_Int_Value) {
                                condInstruction.set_Op(1, new IR_Int_Value(((IR_Int_Value) val).get_Int_Value(), IR_Int32.IR_Int32()));
                            }
                            else{
                                condInstruction.set_Op(1, new IR_Float_Value(((IR_Float_Value) val).get_Float_Value(), IR_Float.IR_Float()));
                            }
                        }
                        if(constmap.containsKey(condInstruction.get_op_index(2))){
                            IR_Value val = constmap.get(condInstruction.get_op_index(2));
                            //System.out.println(val);
                            if (val instanceof IR_Int_Value) {
                                condInstruction.set_Op(2, new IR_Int_Value(((IR_Int_Value) val).get_Int_Value(), IR_Int32.IR_Int32()));
                            }
                            else{
                                condInstruction.set_Op(2, new IR_Float_Value(((IR_Float_Value) val).get_Float_Value(), IR_Float.IR_Float()));
                            }}
                        if (is_Const(condInstruction.get_op_index(1)) && is_Const(condInstruction.get_op_index(2))) {
                            //优化比较
                            IR_Value value = doCompare(condInstruction);
                            //Instruction move = new Move_Instruction(Arrays.asList(value, condInstruction.get_op_index(0)), condInstruction.get_bb());
                            //将原先的比较准换为 move
                            constmap.put(condInstruction.get_op_index(0), value);
                            //记录常量
                            //insertmap.put(move, condInstruction);
                            //删除比较
                            deletelist.add(condInstruction);
                        }
                    }
                    else if (instruction instanceof Zext_Instruction zextInstruction) {
                        IR_Value usevalue = zextInstruction.get_op_index(1);
                        if (constmap.containsKey(usevalue)) {
                            //转move
                            //Move_Instruction move = new Move_Instruction(Arrays.asList(((Move_Instruction) constmap.get(usevalue)).getSrc(), zextInstruction.get_op_index(0)), zextInstruction.get_bb());
                            //deletelist.add(constmap.get(usevalue));
                            //insertmap.put(move, instruction);
                            Move_Instruction move = new Move_Instruction(Arrays.asList(constmap.get(usevalue), zextInstruction.get_op_index(0)), zextInstruction.get_bb());
                            insertmap.put(move, instruction);
                            deletelist.add(zextInstruction);
                            IR_Value value = constmap.get(usevalue);
                            constmap.put(move.getDst(),value);
                        }
                    }
                    //条件跳转
                    else if(instruction instanceof Branch_Instruction branchInstruction && branchInstruction.br_type()){
                        //看能不能转无条件跳转
                        if(constmap.containsKey(branchInstruction.get_op_index(0))){
                            //删除原先的
                            //deletelist.add(constmap.get(branchInstruction.get_op_index(0)));
                            //Move_Instruction moveInstruction = (Move_Instruction) constmap.get(branchInstruction.get_op_index(0));
                            if(constmap.get(branchInstruction.get_op_index(0)).get_Type() == IR_Int1.IR_Int1()){
                                IR_Value value = constmap.get(branchInstruction.get_op_index(0));
                                IR_BaseBlock block = null;
                                if(value instanceof IR_Int_Value irIntValue && irIntValue.get_Int_Value() != 0){
                                    block = branchInstruction.get_block1();
                                    branchInstruction.get_bb().getNextBlock().remove(branchInstruction.get_block2());
                                    if(deleteblocklist.containsKey(branchInstruction.get_block2())){
                                        int z = deleteblocklist.get(branchInstruction.get_block2());
                                        deleteblocklist.put(branchInstruction.get_block2(),z+1);
                                    }
                                    else{
                                        deleteblocklist.put(branchInstruction.get_block2(),1);
                                    }
                                }
                                else{
                                    block = branchInstruction.get_block2();
                                    branchInstruction.get_bb().getNextBlock().remove(branchInstruction.get_block1());
                                    if(deleteblocklist.containsKey(branchInstruction.get_block1())){
                                        int z = deleteblocklist.get(branchInstruction.get_block1());
                                        deleteblocklist.put(branchInstruction.get_block1(),z+1);
                                    }
                                    else{
                                        deleteblocklist.put(branchInstruction.get_block1(),1);
                                    }
                                }
                                //branchInstruction.get_bb().getNextBlock().remove(block);
                                Branch_Instruction branchInstruction1 = new Branch_Instruction(Arrays.asList(block),branchInstruction.get_bb());
                                insertmap.put(branchInstruction1,branchInstruction);
                                deletelist.add(branchInstruction);
                            }
                        }
                    }
                    else if(instruction instanceof Move_Instruction moveInstruction){
                        //deletelist.add(constmap.get(moveInstruction.getSrc()));
                        if(is_Const(moveInstruction.getSrc())){
                            constmap.put(moveInstruction.getDst(),moveInstruction.getSrc());
                        }
                        if(constmap.containsKey(moveInstruction.getSrc())) {
                            IR_Value value = constmap.get(moveInstruction.getSrc());
                            moveInstruction.setSrc(constmap.get(moveInstruction.getSrc()));
                            constmap.put(moveInstruction.getDst(), value);
                        }
                    }
                    else if(instruction instanceof Calc_Instruction calcInstruction){
                        IR_Value lhs = calcInstruction.get_op_index(1);
                        IR_Value rhs = calcInstruction.get_op_index(2);
                        if(constmap.containsKey(lhs)){
                            calcInstruction.set_Op(1,constmap.get(lhs));
                            lhs = constmap.get(lhs);
                        }
                        if(constmap.containsKey(rhs)){
                            calcInstruction.set_Op(2,constmap.get(rhs));
                            rhs = constmap.get(rhs);
                        }
                        if((is_Const(lhs) || constmap.containsKey(lhs) )&&(is_Const(rhs) || constmap.containsKey(rhs))){
                            if(constmap.containsKey(lhs)){
                                lhs = constmap.get(lhs);
                            }
                            if(constmap.containsKey(rhs)){
                                lhs = constmap.get(rhs);
                            }
                            IR_Value value = docala(calcInstruction.get_op_Type(),lhs,rhs,calcInstruction.get_op_index(0).get_Type());
                            constmap.put(calcInstruction.get_op_index(0),value);
                            deletelist.add(calcInstruction);
                        }
                    }
                    else if(instruction instanceof Store_Instruction storeInstruction){
                        
                    }
                }
            }

            for (Map.Entry<Instruction, Instruction> entry : insertmap.entrySet()) {
                int index = entry.getKey().get_bb().get_Instruction_List().indexOf(entry.getValue());
                entry.getKey().get_bb().get_Instruction_List().add(index, entry.getKey());
                changed = true;
            }

            for (Instruction condInstruction : deletelist) {
                changed = true;
                condInstruction.remove();
            }

            for(Map.Entry<IR_BaseBlock,Integer> entry : deleteblocklist.entrySet()){
                //if(entry.getKey().getPreBlock().size() == entry.getValue()){
                    removeBlock(ir_funcBlock,entry.getKey(),entry.getKey().getPreBlock().size() == entry.getValue());
                    changed = true;
                //}
            }
        }
    }

    private IR_Value doCompare(Cond_Instruction condInstruction){
        IR_Value value = null;
        if(condInstruction.get_op_index(1).get_Type() == IR_Float.IR_Float()) {
            //浮点类型
            float lhs = ((IR_Float_Value)condInstruction.get_op_index(1)).get_Float_Value();
            float rhs = ((IR_Float_Value)condInstruction.get_op_index(2)).get_Float_Value();
            switch (condInstruction.get_cond()){
                case IR_UEQ:
                    if(lhs == rhs){
                        value = new IR_Int_Value(1, IR_Int1.IR_Int1());
                    }
                    else{
                        value = new IR_Int_Value(0, IR_Int1.IR_Int1());
                    }
                    break;
                case IR_UNE:
                    if(lhs != rhs){
                        value = new IR_Int_Value(1, IR_Int1.IR_Int1());
                    }
                    else{
                        value = new IR_Int_Value(0, IR_Int1.IR_Int1());
                    }
                    break;
                case IR_UGT:
                    if(lhs > rhs){
                        value = new IR_Int_Value(1, IR_Int1.IR_Int1());
                    }
                    else{
                        value = new IR_Int_Value(0, IR_Int1.IR_Int1());
                    }
                    break;
                case IR_UGE:
                    if(lhs >= rhs){
                        value = new IR_Int_Value(1, IR_Int1.IR_Int1());
                    }
                    else{
                        value = new IR_Int_Value(0, IR_Int1.IR_Int1());
                    }
                    break;
                case IR_ULT:
                    if(lhs < rhs){
                        value = new IR_Int_Value(1, IR_Int1.IR_Int1());
                    }
                    else{
                        value = new IR_Int_Value(0, IR_Int1.IR_Int1());
                    }
                    break;
                case IR_ULE:
                    if(lhs <= rhs){
                        value = new IR_Int_Value(1, IR_Int1.IR_Int1());
                    }
                    else{
                        value = new IR_Int_Value(0, IR_Int1.IR_Int1());
                    }
                    break;
            }
        }
        else{
            int lhs = ((IR_Int_Value)condInstruction.get_op_index(1)).get_Int_Value();
            int rhs = ((IR_Int_Value)condInstruction.get_op_index(2)).get_Int_Value();
            switch (condInstruction.get_cond()){
                case IR_EQ:
                    if(lhs == rhs){
                        value = new IR_Int_Value(1, IR_Int1.IR_Int1());
                    }
                    else{
                        value = new IR_Int_Value(0, IR_Int1.IR_Int1());
                    }
                    break;
                case IR_NE:
                    if(lhs != rhs){
                        value = new IR_Int_Value(1, IR_Int1.IR_Int1());
                    }
                    else{
                        value = new IR_Int_Value(0, IR_Int1.IR_Int1());
                    }
                    break;
                case IR_GT:
                    if(lhs > rhs){
                        value = new IR_Int_Value(1, IR_Int1.IR_Int1());
                    }
                    else{
                        value = new IR_Int_Value(0, IR_Int1.IR_Int1());
                    }
                    break;
                case IR_GEQ:
                    if(lhs >= rhs){
                        value = new IR_Int_Value(1, IR_Int1.IR_Int1());
                    }
                    else{
                        value = new IR_Int_Value(0, IR_Int1.IR_Int1());
                    }
                    break;
                case IR_LT:
                    if(lhs < rhs){
                        value = new IR_Int_Value(1, IR_Int1.IR_Int1());
                    }
                    else{
                        value = new IR_Int_Value(0, IR_Int1.IR_Int1());
                    }
                    break;
                case IR_LEQ:
                    if(lhs <= rhs){
                        value = new IR_Int_Value(1, IR_Int1.IR_Int1());
                    }
                    else{
                        value = new IR_Int_Value(0, IR_Int1.IR_Int1());
                    }
                    break;
            }
        }
        //System.out.println(value);
        //System.out.println(condInstruction);
        return value;
    }
    //检查值是否为常量
    private boolean is_Const(IR_Value value) {
        return value.get_Const_Type() == IR_Const.IR_ConstInt32 ||
                value.get_Const_Type() == IR_Const.IR_ConstFloat || value.get_Const_Type() == IR_Const.IR_ConstInt1;
    }

    private void removeBlock(IR_FuncBlock funcBlock,IR_BaseBlock block,boolean isdelete){
        //module.buildCFG();
        //System.out.println(block + "  " +isdelete);
        //需要删除这个基本块 先处理掉后继块
        if(isdelete) {
            for(IR_BaseBlock nxtblock:block.getNextBlock()){
                for(Instruction instruction:nxtblock.get_Instruction_List()){
                    if(instruction instanceof Phi_Instruction phiInstruction){
                        phiInstruction.remove_ComingValues(block);
                    }
                    else{
                        break;
                    }
                }
                nxtblock.getPreBlock().remove(block);
            }
            funcBlock.get_bb_List().remove(block);
        }
        else {
            //此时还需要注意当前 bb 的phi指令
            IR_BaseBlock bb = null;
            for (IR_BaseBlock preblock : block.getPreBlock()) {
                //条件跳转转无条件跳转
                if (preblock.getNextBlock().contains(block)) {
                    continue;
                }
                bb = preblock;
                break;
            }
            for (Instruction instruction : block.get_Instruction_List()) {
                if (instruction instanceof Phi_Instruction phiInstruction) {
                    phiInstruction.remove_ComingValues(bb);
                } else {
                    break;
                }
            }
            block.getPreBlock().remove(bb);
        }
        //module
    }

    private IR_Value docala(String type, IR_Value lhs, IR_Value rhs, IR_Type res){
        IR_Value value = null;
        if(res == IR_Float.IR_Float()){
            float l = ((IR_Float_Value) lhs).get_Float_Value();
            float r = ((IR_Float_Value) rhs).get_Float_Value();
            switch(type){
                case FADD:
                    value = new IR_Float_Value(l + r);
                    break;
                case FSUB:
                    value = new IR_Float_Value(l - r);
                    break;
                case FMUL:
                    value = new IR_Float_Value(l * r);
                    break;
                case FDIV:
                    value = new IR_Float_Value(l / r);
                    break;
            }

        }
        else{
            int l = ((IR_Int_Value) lhs).get_Int_Value();
            int r = ((IR_Int_Value) rhs).get_Int_Value();
            switch (type) {
                case XOR:
                    value = new IR_Int_Value(l ^ r, res);
                    break;
                case ADD:
                    value = new IR_Int_Value(l + r, res);
                    break;
                case SUB:
                    value = new IR_Int_Value(l - r, res);
                    break;
                case MUL:
                    value = new IR_Int_Value(l * r, res);
                    break;
                case SDIV:
                    value = new IR_Int_Value(l / r, res);
                    break;
                case SHL:
                    value = new IR_Int_Value(l << r, res); // 左移操作
                    break;
                case ASHR:
                    value = new IR_Int_Value(l >> r, res); // 算术右移操作
                    break;
                case LSHR:
                    value = new IR_Int_Value(l >>> r, res); // 逻辑右移操作
                    break;
                case SREM:
                    value = new IR_Int_Value(l % r, res); // 取模操作，除数不能为零
                    break;
            }
        }

        return value;
    }
}
